{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  import pydrake\n",
    "  import underactuated\n",
    "except ImportError:\n",
    "  !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "  from jupyter_setup import setup_underactuated\n",
    "  setup_underactuated()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trajectory optimization for the double integrator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pydrake.all import eq, MathematicalProgram, Solve, Variable\n",
    "\n",
    "# Discrete-time approximation of the double integrator.\n",
    "dt = 0.01;\n",
    "A = np.eye(2) + dt*np.mat('0 1; 0 0')\n",
    "B = dt*np.mat('0; 1')\n",
    "\n",
    "prog = MathematicalProgram()\n",
    "\n",
    "N = 284  # Note: I had to do a manual \"line search\" to find this.\n",
    "\n",
    "# Create decision variables\n",
    "u = np.empty((1, N-1), dtype=Variable)\n",
    "x = np.empty((2, N), dtype=Variable)\n",
    "for n in range(N-1):\n",
    "  u[:,n] = prog.NewContinuousVariables(1, 'u' + str(n))\n",
    "  x[:,n] = prog.NewContinuousVariables(2, 'x' + str(n))\n",
    "x[:,N-1] = prog.NewContinuousVariables(2, 'x' + str(N))\n",
    "\n",
    "# Add constraints\n",
    "x0 = [-2, 0]\n",
    "prog.AddBoundingBoxConstraint(x0, x0, x[:,0])\n",
    "for n in range(N-1):\n",
    "  # Will eventually be prog.AddConstraint(x[:,n+1] == A@x[:,n] + B@u[:,n])\n",
    "  # See drake issues 12841 and 8315\n",
    "  prog.AddConstraint(eq(x[:,n+1],A.dot(x[:,n]) + B.dot(u[:,n])))\n",
    "  prog.AddBoundingBoxConstraint(-1, 1, u[:,n])\n",
    "xf = [0, 0]\n",
    "prog.AddBoundingBoxConstraint(xf, xf, x[:, N-1])\n",
    "\n",
    "result = Solve(prog)\n",
    "\n",
    "x_sol = result.GetSolution(x)\n",
    "assert(result.is_success()), \"Optimization failed\"\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(x_sol[0,:], x_sol[1,:])\n",
    "plt.xlabel('q')\n",
    "plt.ylabel('qdot');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trajectory optimization using the DirectTranscription class\n",
    "\n",
    "Because this pattern of making decision variables that are indexed over time, adding the dynamic constraints, defining the running cost and constraints, is so common, we have wrappers in drake on top of `MathematicalProgram` which handle these details for you.\n",
    "\n",
    "The optimization below is identical to the example above, but using this helper class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pydrake.all import LinearSystem, DirectTranscription\n",
    "\n",
    "# Discrete-time approximation of the double integrator.\n",
    "dt = 0.01;\n",
    "A = np.eye(2) + dt*np.mat('0 1; 0 0')\n",
    "B = dt*np.mat('0; 1')\n",
    "C = np.eye(2)\n",
    "D = np.zeros((2,1))\n",
    "sys = LinearSystem(A, B, C, D, dt)\n",
    "\n",
    "prog = DirectTranscription(sys, sys.CreateDefaultContext(), N)\n",
    "prog.AddBoundingBoxConstraint(x0, x0, prog.initial_state())\n",
    "prog.AddBoundingBoxConstraint(xf, xf, prog.final_state())\n",
    "prog.AddConstraintToAllKnotPoints(prog.input()[0] <= 1)\n",
    "prog.AddConstraintToAllKnotPoints(prog.input()[0] >= -1)\n",
    "\n",
    "result = Solve(prog)\n",
    "x_sol = prog.ReconstructStateTrajectory(result)\n",
    "assert(result.is_success()), \"Optimization failed\"\n",
    "\n",
    "plt.figure()\n",
    "x_values = x_sol.vector_values(x_sol.get_segment_times())\n",
    "\n",
    "plt.plot(x_values[0,:], x_values[1,:])\n",
    "plt.xlabel('q')\n",
    "plt.ylabel('qdot');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One thing that I'm very proud of (it was a lot of work!) is the fact that drake is often smart enough to introspect your system, costs, and constraints and understand whether you have formulated a convex problem or a non-convex one.  The optimization above calls a convex optimization solver.  But if you had passed in a nonlinear system instead, it would have switched to calling a solver that supports nonlinear programming."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Direct Collocation for the Pendulum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from underactuated.jupyter import SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "from pydrake.examples.pendulum import PendulumPlant, PendulumState\n",
    "from pydrake.all import DirectCollocation, PiecewisePolynomial, Solve\n",
    "from underactuated.pendulum import PendulumVisualizer\n",
    "\n",
    "\n",
    "plant = PendulumPlant()\n",
    "context = plant.CreateDefaultContext()\n",
    "\n",
    "N = 21\n",
    "max_dt = 0.5\n",
    "max_tf = N * max_dt\n",
    "dircol = DirectCollocation(plant,\n",
    "                           context,\n",
    "                           num_time_samples=N,\n",
    "                           minimum_timestep=0.05,\n",
    "                           maximum_timestep=max_dt)\n",
    "\n",
    "dircol.AddEqualTimeIntervalsConstraints()\n",
    "\n",
    "torque_limit = 3.0  # N*m.\n",
    "u = dircol.input()\n",
    "dircol.AddConstraintToAllKnotPoints(-torque_limit <= u[0])\n",
    "dircol.AddConstraintToAllKnotPoints(u[0] <= torque_limit)\n",
    "\n",
    "initial_state = PendulumState()\n",
    "initial_state.set_theta(0.0)\n",
    "initial_state.set_thetadot(0.0)\n",
    "dircol.AddBoundingBoxConstraint(initial_state.get_value(),\n",
    "                                initial_state.get_value(),\n",
    "                                dircol.initial_state())\n",
    "# More elegant version is blocked on drake #8315:\n",
    "# dircol.AddLinearConstraint(\n",
    "#     dircol.initial_state() == initial_state.get_value())\n",
    "\n",
    "final_state = PendulumState()\n",
    "final_state.set_theta(math.pi)\n",
    "final_state.set_thetadot(0.0)\n",
    "dircol.AddBoundingBoxConstraint(final_state.get_value(),\n",
    "                                final_state.get_value(), dircol.final_state())\n",
    "# dircol.AddLinearConstraint(dircol.final_state() == final_state.get_value())\n",
    "\n",
    "R = 10  # Cost on input \"effort\".\n",
    "dircol.AddRunningCost(R * u[0]**2)\n",
    "\n",
    "initial_x_trajectory = PiecewisePolynomial.FirstOrderHold(\n",
    "    [0., 4.], [initial_state.get_value(),\n",
    "               final_state.get_value()])\n",
    "dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)\n",
    "\n",
    "result = Solve(dircol)\n",
    "assert result.is_success()\n",
    "\n",
    "x_trajectory = dircol.ReconstructStateTrajectory(result)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "x_knots = np.hstack([\n",
    "    x_trajectory.value(t) for t in np.linspace(x_trajectory.start_time(),\n",
    "                                               x_trajectory.end_time(), 100)\n",
    "])\n",
    "\n",
    "ax.plot(x_knots[0, :], x_knots[1, :]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "\n",
    "# Animate the result.\n",
    "vis = PendulumVisualizer(show=False)\n",
    "ani = vis.animate(x_trajectory)\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Direct Collocation for the Acrobot\n",
    "\n",
    "Almost identical code works to swing-up the Acrobot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pydrake.all import (DirectCollocation, DiagramBuilder, PiecewisePolynomial,\n",
    "                         PlanarSceneGraphVisualizer, SceneGraph, Simulator,\n",
    "                         Solve, TrajectorySource)\n",
    "from pydrake.examples.acrobot import AcrobotGeometry, AcrobotPlant\n",
    "from underactuated import FindResource\n",
    "\n",
    "plant = AcrobotPlant()\n",
    "context = plant.CreateDefaultContext()\n",
    "\n",
    "dircol = DirectCollocation(plant,\n",
    "                           context,\n",
    "                           num_time_samples=21,\n",
    "                           minimum_timestep=0.05,\n",
    "                           maximum_timestep=0.2)\n",
    "\n",
    "dircol.AddEqualTimeIntervalsConstraints()\n",
    "\n",
    "# Add input limits.\n",
    "torque_limit = 8.0  # N*m.\n",
    "u = dircol.input()\n",
    "dircol.AddConstraintToAllKnotPoints(-torque_limit <= u[0])\n",
    "dircol.AddConstraintToAllKnotPoints(u[0] <= torque_limit)\n",
    "\n",
    "initial_state = (0., 0., 0., 0.)\n",
    "dircol.AddBoundingBoxConstraint(initial_state, initial_state,\n",
    "                                dircol.initial_state())\n",
    "# More elegant version is blocked on drake #8315:\n",
    "# dircol.AddLinearConstraint(dircol.initial_state() == initial_state)\n",
    "\n",
    "final_state = (math.pi, 0., 0., 0.)\n",
    "dircol.AddBoundingBoxConstraint(final_state, final_state, dircol.final_state())\n",
    "# dircol.AddLinearConstraint(dircol.final_state() == final_state)\n",
    "\n",
    "R = 10  # Cost on input \"effort\".\n",
    "dircol.AddRunningCost(R * u[0]**2)\n",
    "\n",
    "# Add a final cost equal to the total duration.\n",
    "dircol.AddFinalCost(dircol.time())\n",
    "\n",
    "initial_x_trajectory = PiecewisePolynomial.FirstOrderHold(\n",
    "    [0., 4.], np.column_stack((initial_state, final_state)))  # yapf: disable\n",
    "dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)\n",
    "\n",
    "result = Solve(dircol)\n",
    "assert (result.is_success())\n",
    "\n",
    "u_trajectory = dircol.ReconstructInputTrajectory(result)\n",
    "times = np.linspace(u_trajectory.start_time(), u_trajectory.end_time(), 100)\n",
    "u_lookup = np.vectorize(u_trajectory.value)\n",
    "u_values = u_lookup(times)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(times, u_values)\n",
    "plt.xlabel(\"time (seconds)\")\n",
    "plt.ylabel(\"force (Newtons)\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "x_trajectory = dircol.ReconstructStateTrajectory(result)\n",
    "\n",
    "builder = DiagramBuilder()\n",
    "source = builder.AddSystem(TrajectorySource(x_trajectory))\n",
    "scene_graph = builder.AddSystem(SceneGraph())\n",
    "AcrobotGeometry.AddToBuilder(builder, source.get_output_port(0), scene_graph)\n",
    "visualizer = builder.AddSystem(\n",
    "    PlanarSceneGraphVisualizer(scene_graph, xlim=[-4., 4.], ylim=[-4., 4.], \n",
    "                               show=plt_is_interactive))\n",
    "builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
    "                visualizer.get_input_port(0))\n",
    "simulator = Simulator(builder.Build())\n",
    "\n",
    "# Simulate and animate\n",
    "AdvanceToAndVisualize(simulator, visualizer, x_trajectory.end_time())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trajectory stabilization with (finite-horizon) LQR\n",
    "\n",
    "The visualization above was an animation of the solution trajectory, it was not a simulation!  In fact, if you simulate only the planned torques connected to the Acrobot, then you will see that the trajectory does not follow the planned state trajectory.  The trajectory optimization only satisfies the dynamic constraints to some loose tolerance, and the (open-loop) trajectory is unstable.\n",
    "\n",
    "To do better, we must stabilizing the planned trajectory with feedback.  We do that here with a finite-horizon LQR design, which takes a time-varying linearization of the plant along the trajectory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.all import MakeFiniteHorizonLinearQuadraticRegulator, FiniteHorizonLinearQuadraticRegulatorOptions, LogOutput\n",
    "\n",
    "options = FiniteHorizonLinearQuadraticRegulatorOptions()\n",
    "options.x0 = dircol.ReconstructStateTrajectory(result)\n",
    "options.u0 = dircol.ReconstructInputTrajectory(result)\n",
    "\n",
    "builder = DiagramBuilder()\n",
    "plant = builder.AddSystem(AcrobotPlant())\n",
    "context = plant.CreateDefaultContext()\n",
    "Q = np.diag([10.,10.,1.,1.])\n",
    "options.Qf = Q\n",
    "regulator = builder.AddSystem(MakeFiniteHorizonLinearQuadraticRegulator(\n",
    "  plant, context, t0=options.u0.start_time(), \n",
    "  tf=options.u0.end_time(), Q=Q, R=np.eye(1),\n",
    "  options=options))\n",
    "builder.Connect(regulator.get_output_port(0), plant.get_input_port(0))\n",
    "builder.Connect(plant.get_output_port(0), regulator.get_input_port(0))\n",
    "input_logger = LogOutput(regulator.get_output_port(0), builder)\n",
    "state_logger = LogOutput(plant.get_output_port(0), builder)\n",
    "\n",
    "scene_graph = builder.AddSystem(SceneGraph())\n",
    "AcrobotGeometry.AddToBuilder(builder, plant.get_output_port(0), scene_graph)\n",
    "visualizer = builder.AddSystem(\n",
    "    PlanarSceneGraphVisualizer(scene_graph, xlim=[-4., 4.], ylim=[-4., 4.], \n",
    "                               show=plt_is_interactive))\n",
    "builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
    "                visualizer.get_input_port(0))\n",
    "diagram = builder.Build()\n",
    "simulator = Simulator(diagram)\n",
    "\n",
    "# Simulate and animate\n",
    "AdvanceToAndVisualize(simulator, visualizer, options.u0.end_time())\n",
    "\n",
    "fig, ax = plt.subplots(2,1)\n",
    "ax[0].plot(options.u0.get_segment_times(), options.u0.vector_values(options.u0.get_segment_times()).T)\n",
    "ax[0].plot(input_logger.sample_times(), input_logger.data().T)\n",
    "ax[0].legend(('u nominal','u actual'))\n",
    "\n",
    "ax[1].plot(options.x0.get_segment_times(), options.x0.vector_values(options.x0.get_segment_times()).T, 'b')\n",
    "ax[1].plot(state_logger.sample_times(), state_logger.data().T, 'g')\n",
    "ax[1].legend(('x nominal','x actual'));\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Direct Collocation for the Cart-Pole\n",
    "\n",
    "While the previous two examples used equations of motion that we derived and typed in manually, for the Cart-Pole we will use the Drake dynamics engine and visualization via [`MultibodyPlant`](https://drake.mit.edu/doxygen_cxx/classdrake_1_1multibody_1_1_multibody_plant.html) and [`SceneGraph`](https://drake.mit.edu/doxygen_cxx/classdrake_1_1geometry_1_1_scene_graph.html).  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pydrake.all import (DiagramBuilder, DirectCollocation, MultibodyPlant,\n",
    "                         MultibodyPositionToGeometryPose, Parser,\n",
    "                         PiecewisePolynomial, PlanarSceneGraphVisualizer,\n",
    "                         SceneGraph, Simulator, Solve, TrajectorySource)\n",
    "from underactuated import FindResource\n",
    "\n",
    "plant = MultibodyPlant(time_step=0.0)\n",
    "scene_graph = SceneGraph()\n",
    "plant.RegisterAsSourceForSceneGraph(scene_graph)\n",
    "file_name = FindResource(\"models/cartpole.urdf\")\n",
    "Parser(plant).AddModelFromFile(file_name)\n",
    "plant.Finalize()\n",
    "\n",
    "context = plant.CreateDefaultContext()\n",
    "dircol = DirectCollocation(\n",
    "    plant,\n",
    "    context,\n",
    "    num_time_samples=21,\n",
    "    minimum_timestep=0.1,\n",
    "    maximum_timestep=0.4,\n",
    "    input_port_index=plant.get_actuation_input_port().get_index())\n",
    "\n",
    "dircol.AddEqualTimeIntervalsConstraints()\n",
    "\n",
    "initial_state = (0., 0., 0., 0.)\n",
    "dircol.AddBoundingBoxConstraint(initial_state, initial_state,\n",
    "                                dircol.initial_state())\n",
    "# More elegant version is blocked by drake #8315:\n",
    "# dircol.AddLinearConstraint(dircol.initial_state() == initial_state)\n",
    "\n",
    "final_state = (0., math.pi, 0., 0.)\n",
    "dircol.AddBoundingBoxConstraint(final_state, final_state, dircol.final_state())\n",
    "# dircol.AddLinearConstraint(dircol.final_state() == final_state)\n",
    "\n",
    "R = 10  # Cost on input \"effort\".\n",
    "u = dircol.input()\n",
    "dircol.AddRunningCost(R * u[0]**2)\n",
    "\n",
    "# Add a final cost equal to the total duration.\n",
    "dircol.AddFinalCost(dircol.time())\n",
    "\n",
    "initial_x_trajectory = PiecewisePolynomial.FirstOrderHold(\n",
    "    [0., 4.], np.column_stack((initial_state, final_state)))  # yapf: disable\n",
    "dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)\n",
    "\n",
    "result = Solve(dircol)\n",
    "assert result.is_success()\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "u_trajectory = dircol.ReconstructInputTrajectory(result)\n",
    "times = np.linspace(u_trajectory.start_time(), u_trajectory.end_time(), 100)\n",
    "u_lookup = np.vectorize(u_trajectory.value)\n",
    "u_values = u_lookup(times)\n",
    "\n",
    "ax.plot(times, u_values)\n",
    "ax.set_xlabel(\"time (seconds)\")\n",
    "ax.set_ylabel(\"force (Newtons)\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Animate the results.\n",
    "from IPython.display import HTML\n",
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "x_trajectory = dircol.ReconstructStateTrajectory(result)\n",
    "\n",
    "# TODO(russt): Add some helper methods to make this workflow cleaner.\n",
    "builder = DiagramBuilder()\n",
    "source = builder.AddSystem(TrajectorySource(x_trajectory))\n",
    "builder.AddSystem(scene_graph)\n",
    "pos_to_pose = builder.AddSystem(\n",
    "    MultibodyPositionToGeometryPose(plant, input_multibody_state=True))\n",
    "builder.Connect(source.get_output_port(0), pos_to_pose.get_input_port())\n",
    "builder.Connect(pos_to_pose.get_output_port(),\n",
    "                scene_graph.get_source_pose_port(plant.get_source_id()))\n",
    "\n",
    "visualizer = builder.AddSystem(\n",
    "    PlanarSceneGraphVisualizer(scene_graph,\n",
    "                               xlim=[-2, 2],\n",
    "                               ylim=[-1.25, 2],\n",
    "                               show=plt_is_interactive))\n",
    "builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
    "                visualizer.get_input_port(0))\n",
    "simulator = Simulator(builder.Build())\n",
    "\n",
    "AdvanceToAndVisualize(simulator, visualizer, x_trajectory.end_time())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
